/*********************************************************************************
*   myecho.cpp                    
*     
*   Alberto De Pena
*   February 9, 2011
*                                 
*   Echo client - version 1      
**********************************************************************************/
#include <winsock2.h>
#include <iostream>

using namespace std;

#pragma comment (lib, "ws2_32.lib")

const int  MAXBUF  = 128;
const int  MAXNAME = 80;

SOCKET connectsock (char*, char*, char*);

int main ( int argc, char* argv[] ) 
{
    char msg[MAXBUF] = "";
    char buf[MAXBUF] = "";
    char host[MAXNAME] = "";
    char service[MAXNAME] = "";
    int  len;
    SOCKET s;

    /* Must always initialize the Windows Sockets TCP/IP library */
    WORD wVersion = 0x0202;
    WSADATA wsaData;
    int iResult = WSAStartup( wVersion, &wsaData);
    if (iResult != 0)
    {
        cout << "Unable to initialize Windows Socket library." << endl;
        return 0;
    }

    if( argc == 1 )
    {
        cout << "Usage: myecho server_address [port #]" << endl;
        return 0;
    }

    if ( argc == 2 )
    {
        strcpy_s(host, sizeof(host), argv[1] );    
        strcpy_s(service, sizeof(service), "7");
    }
    else if( argc == 3)
    {
        strcpy_s(host, sizeof(host), argv[1] );    
        strcpy_s(service, sizeof(service), argv[2] );
    }

    cout << "Message? ";
    cin.getline(msg, MAXBUF);
    while( stricmp( msg, "QUIT") )
    {
        s = connectsock(host, service, "TCP");
        if (s != INVALID_SOCKET) 
        {
            len = send(s, msg, (int) strlen(msg), 0);
            if (len > 0) 
            {
                cout << "Number of bytes sent: " << len << endl;
                cout << endl;
            } else
            {
                cout << "Unable to send message to remote machine." << endl;
                return 0;
            }

            // Message was sent.  Is there a reply from the remote machine?
            len = recv(s, buf, sizeof(buf), 0);
            if (len > 0) 
            {
                cout << "Received message: ";
                for (int i=0; i<len; i++) cout << buf[i];
                cout << endl;
                cout << "Number of bytes received: " << len << endl;
            }
            cout << endl;

            /* Close the socket (and thus the connection) */
            shutdown(s,1);
            closesocket(s);
        }
        cout << "Message? ";
        cin.getline(msg, MAXBUF);
    }
        WSACleanup();
}


/*------------------------------------------------------------------
 * connectsock - allocate & connect a remote socket using TCP or UDP
 *------------------------------------------------------------------
 */
SOCKET connectsock (char *host, char *service, char *protocol)
{   // Parameter List:
    // [IN] host    : String containing remote host's IP address
    // [IN] service : String containing remote port number (eg. "7")
    // [IN] protocol: String containing protocol type (either "tcp" or "udp")
    // Return value : SOCKET handle of connected socket.

    //    NOTE: This function requires winsock2.h header file and ws2_32.lib library file.

    struct sockaddr_in remote; /* an Internet endpoint address structure  */
    SOCKET s;                   /* socket descriptor                       */

    /* Allocate and connect socket - type depends upon protocol */
    if (_stricmp(protocol, "tcp") == 0)  
    {   
        // if TCP request, do following
        s = socket(AF_INET, SOCK_STREAM, 0);
        if (s < 0 || s == INVALID_SOCKET) 
        {
            cout << "Cannot create socket" << endl;
            return INVALID_SOCKET;
        }

        // Zero the remote struct
        memset(&remote, 0, sizeof(remote));

        // Set the address family to AF_INET (that is, "Internet")
        remote.sin_family = AF_INET;

        // Set the port number (this version assumes service contains string of digits)
        remote.sin_port = htons( (u_short) atoi(service));
        struct servent *pse;
        if (remote.sin_port == 0)
        {                       
            // if that doesn't work, call service function
            pse = getservbyname(service,NULL);
            if (pse) 
            {
                remote.sin_port = pse->s_port;
            }
            else
            {
                cout << "Invalid service request." << endl;
                return INVALID_SOCKET;
            }
        }

        // Set remote IP address (this version assumes host contains a string in dotted decimal format)
        remote.sin_addr.s_addr = inet_addr(host);
        struct hostent *phe;
        if (remote.sin_addr.s_addr == INADDR_NONE)
        {   
            // If that doesn't work, send the hostname to DNS
            phe = gethostbyname(host);
            if (phe)
            {
                memmove(&remote.sin_addr.s_addr, phe->h_addr, phe->h_length);
            }
            else
            {
                cout << "Could not find host." << endl;
                return INVALID_SOCKET;
            }
        }

        /* Connect the socket */
        int status = connect(s, (LPSOCKADDR) &remote, sizeof(SOCKADDR));
        if (status == SOCKET_ERROR) 
        {
            cout << "Remote host/service not found - or connection refused" << endl;
            return INVALID_SOCKET;
        }
    }
    else if (_stricmp(protocol,"udp") == 0) 
    {   
        // if UDP request, do following
        s = socket(AF_INET, SOCK_DGRAM,  0);
        if (s < 0 || s == INVALID_SOCKET) 
        {
            cout << "Cannot create socket" << endl;
            return INVALID_SOCKET;
        }
    }
    else 
    {   
        // Unknown or unsupported protocol request
        cout << "Invalid Protocol" << endl;
        return INVALID_SOCKET;
    }
        
    return s;
}

